use crate::utils::system_allocator::BaseSystemAllocator;
use std::alloc::Layout;
use std::ops::{Deref, DerefMut};
use std::ptr::NonNull;

type OsAllocator = BaseSystemAllocator;

/// OsBox 是对手动进行内存管理的基础封装，功能和标准库的 Box 类似。
///
///
/// 使用 OsBox 管理的内部结构必须实现 OsBoxBody，
/// - OsBoxBody 允许用户自定义初始化和销毁过程
/// - OsBox 保证 OsBoxBody 的 init 和 before_drop 函数一定被调用
///
///
/// OsBox 允许通过 new_with_layout 函数，从外部接受一个 layout，
/// 从而创建一个实际物理空间比 T 所需空间更大的 OsBox。
/// 这个功能在创建 Segment 时很有用。
///
/// Segment 和 Heap 都是用过 OsBox 创建的。
pub(crate) struct OsBox<T>
where
    T: OsBoxBody,
{
    ptr: NonNull<T>,
    layout: Layout,
    is_leaked: bool,
}

impl<T> OsBox<T>
where
    T: OsBoxBody,
{
    pub(crate) fn new() -> Option<Self> {
        Self::new_with_layout(Layout::new::<T>())
    }

    pub(crate) fn new_with_layout(layout: Layout) -> Option<Self> {
        debug_assert!(layout.size() >= Layout::new::<T>().size());
        unsafe {
            let ptr = OsAllocator {}.alloc(layout) as *mut T;
            if ptr.is_null() {
                return None;
            }
            ptr.write(T::init());
            Some(Self {
                ptr: NonNull::new_unchecked(ptr),
                layout,
                is_leaked: false,
            })
        }
    }

    pub(crate) fn as_ptr(&self) -> NonNull<T> {
        self.ptr
    }

    #[allow(dead_code)]
    pub(crate) fn leak(mut self) {
        self.is_leaked = true;
    }
    #[allow(dead_code)]
    pub(crate) unsafe fn leak_box(ptr: NonNull<T>) -> Self {
        Self {
            ptr,
            layout: Layout::new::<T>(),
            is_leaked: true,
        }
    }
}

impl<T> AsMut<T> for OsBox<T>
where
    T: OsBoxBody,
{
    fn as_mut(&mut self) -> &mut T {
        unsafe { self.ptr.as_mut() }
    }
}

impl<T> AsRef<T> for OsBox<T>
where
    T: OsBoxBody,
{
    fn as_ref(&self) -> &T {
        unsafe { self.ptr.as_ref() }
    }
}

impl<T> Deref for OsBox<T>
where
    T: OsBoxBody,
{
    type Target = T;

    fn deref(&self) -> &Self::Target {
        self.as_ref()
    }
}

impl<T> DerefMut for OsBox<T>
where
    T: OsBoxBody,
{
    fn deref_mut(&mut self) -> &mut Self::Target {
        self.as_mut()
    }
}

pub(crate) trait OsBoxBody {
    fn init() -> Self;
    fn before_drop(&mut self);
}

impl<T> Drop for OsBox<T>
where
    T: OsBoxBody,
{
    fn drop(&mut self) {
        if self.is_leaked {
            return;
        }
        unsafe {
            self.ptr.as_mut().before_drop();
            let heap_ptr = self.ptr.as_ptr() as *mut u8;
            OsAllocator {}.dealloc(heap_ptr, self.layout)
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::utils::os_box::{OsBox, OsBoxBody};
    use std::sync::atomic::{AtomicUsize, Ordering};

    static CALL_CONST: AtomicUsize = AtomicUsize::new(0);
    struct TestStruct {
        a: u32,
        b: u8,
    }

    impl OsBoxBody for TestStruct {
        fn init() -> Self {
            CALL_CONST.fetch_add(1, Ordering::SeqCst);
            Self { a: 10, b: 20 }
        }

        fn before_drop(&mut self) {
            assert_eq!(self.a, 10);
            assert_eq!(self.b, 20);
            CALL_CONST.fetch_add(1, Ordering::SeqCst);
        }
    }
    #[test]
    fn test_os_box() {
        CALL_CONST.store(0, Ordering::Relaxed);
        assert_eq!(CALL_CONST.load(Ordering::Relaxed), 0);
        let t = OsBox::<TestStruct>::new();
        assert_eq!(CALL_CONST.load(Ordering::Relaxed), 1);
        drop(t);
        assert_eq!(CALL_CONST.load(Ordering::Relaxed), 2);
    }
}
